# coding=utf-8
##
## Author: jmdvirus@aliyun.com
##
## Create: 2019年02月16日 星期六 10时25分09秒
##

import numpy as np
import cv2
import matplotlib.pyplot as plt

np.random.seed(42)

def generate_data(num_samples, num_features=2):
    data_size = (num_samples, num_features)
    data = np.random.randint(0, 100, size=data_size)

    label_size = (num_samples, 1)
    labels = np.random.randint(0, 2, size=label_size)

    return data.astype(np.float32), labels

def plot_data(all_blue, all_red):
    plt.scatter(all_blue[:, 0], all_blue[:, 1], c='b',
    marker='s', s=180)
    plt.scatter(all_red[:, 0], all_red[:, 1], c='r',
    marker='^', s=180)

    plt.xlabel('x (feature1)')
    plt.ylabel('y (feature2)')

def train(train_data, labels, newdata):
    knn = cv2.ml.KNearest_create()
    knn.train(train_data, cv2.ml.ROW_SAMPLE, labels)

    ret, results, neighbor, dist = knn.findNearest(newdata, 7)
    print("ret: ", ret)
    print("results: ", results)
    print("neightbor: ", neighbor)
    print("dist: ", dist)

if __name__ == "__main__":
    train_data, labels = generate_data(11)
    print("train_data: ", train_data)
    print("labels: " , labels)

    blue = train_data[labels.ravel() == 0]
    red = train_data[labels.ravel() == 1]

    print("blue: ", blue)

    plot_data(blue, red)

    newcomer, _ = generate_data(1)
    plt.plot(newcomer[0,0], newcomer[0, 1], 'go', markersize=14)

    train(train_data, labels, newcomer)

    plt.show()
